package org.witsoft.demo.graphql.extend.execution;

import graphql.ExecutionResult;
import graphql.ExecutionResultImpl;
import graphql.execution.*;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.function.BiConsumer;

@Component
public class CustomerExecutionStrategy extends AsyncExecutionStrategy{

    public static final ThreadLocal<Set<String>> REMOVE_FIELDS_HOLDER = new ThreadLocal<Set<String>>(){
        @Override
        protected Set<String> initialValue() {
            return new HashSet<>();
        }
    };

    @Override
    public CompletableFuture<ExecutionResult> execute(ExecutionContext executionContext, ExecutionStrategyParameters parameters) throws NonNullableFieldWasNullException {
        REMOVE_FIELDS_HOLDER.get().clear();
        CompletableFuture<ExecutionResult> result = super.execute(executionContext, parameters);
        return result;
    }

    @Override
    protected BiConsumer<List<ExecutionResult>, Throwable> handleResults(ExecutionContext executionContext, List<String> fieldNames, CompletableFuture<ExecutionResult> overallResult) {
        return (List<ExecutionResult> results, Throwable exception) -> {
            if (exception != null) {
                handleNonNullException(executionContext, overallResult, exception);
                return;
            }
            Map<String, Object> resolvedValuesByField = new LinkedHashMap<>(fieldNames.size());
            int ix = 0;
            for (ExecutionResult executionResult : results) {
                String fieldName = fieldNames.get(ix++);
                Object data = executionResult.getData();
                if(REMOVE_FIELDS_HOLDER.get().contains(fieldName)){
                    continue;
                }
                resolvedValuesByField.put(fieldName, data);
            }
            overallResult.complete(new ExecutionResultImpl(resolvedValuesByField, executionContext.getErrors()));
        };
    }
}
